import streamlit as st
import akshare as ak
import pandas as pd
import numpy as np
from datetime import datetime
import matplotlib.pyplot as plt
import seaborn as sns

# 设置中文显示
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False

@st.cache_data
def get_index_data():
    """获取上证指数历史数据"""
    try:
        df = ak.stock_zh_index_daily(symbol="sh000001")
        df['date'] = pd.to_datetime(df['date'])
        df['year'] = df['date'].dt.year
        df['month'] = df['date'].dt.month
        df['day'] = df['date'].dt.day
        return df.sort_values('date')
    except Exception as e:
        st.error(f"数据获取失败: {str(e)}")
        return pd.DataFrame()

def analyze_monthly_data(df, selected_year, selected_month):
    """分析指定月份的历史数据"""
    start_year = selected_year - 10
    results = []
    
    for year in range(start_year, selected_year + 1):
        month_data = df[(df['year'] == year) & (df['month'] == selected_month)]
        
        if len(month_data) >= 10:  # 至少有10个交易日
            # 计算全月涨跌幅
            monthly_change = (month_data.iloc[-1]['close'] / month_data.iloc[0]['close'] - 1) * 100
            
            # 计算上半月涨跌幅（1-15日）
            first_half = month_data[month_data['day'] <= 15]
            if len(first_half) > 0:
                first_half_change = (first_half.iloc[-1]['close'] / first_half.iloc[0]['close'] - 1) * 100
            else:
                first_half_change = np.nan
                
            # 计算下半月涨跌幅（16日-月末）
            second_half = month_data[month_data['day'] > 15]
            if len(second_half) > 0:
                second_half_change = (second_half.iloc[-1]['close'] / second_half.iloc[0]['close'] - 1) * 100
            else:
                second_half_change = np.nan
                
            results.append({
                '年份': year,
                '全月涨跌幅 (%)': monthly_change,
                '上半月涨跌幅 (%)': first_half_change,
                '下半月涨跌幅 (%)': second_half_change,
                '交易天数': len(month_data)
            })
    
    return pd.DataFrame(results)

def display_analysis(df):
    """显示分析界面"""
    current_year = datetime.now().year
    current_month = datetime.now().month
    
    st.title("上证指数月度历史分析")
    
    # 创建并排的选择器
    col1, col2 = st.columns(2)
    with col1:
        selected_year = st.slider(
            "选择结束年份",
            min_value=2015,
            max_value=current_year,
            value=current_year,
            step=1
        )
    with col2:
        selected_month = st.select_slider(
            "选择分析月份",
            options=list(range(1, 13)),
            value=current_month
        )
    
    st.markdown(f"### {selected_year-10}-{selected_year}年 {selected_month}月历史表现")
    
    # 分析数据
    results_df = analyze_monthly_data(df, selected_year, selected_month)
    
    if results_df.empty:
        st.warning("所选时间段数据不足，请调整年份范围")
        return
    
    # 格式化显示表格
    def color_cell(val):
        if isinstance(val, (int, float)):
            color = '#FF3333' if val > 0 else '#4CB050'
            return f'background-color: {color}; color: white; font-weight: bold'
        return ''
    
    styled_df = (
        results_df.style
        .format({
            '全月涨跌幅 (%)': '{:.2f}%',
            '上半月涨跌幅 (%)': '{:.2f}%',
            '下半月涨跌幅 (%)': '{:.2f}%'
        })
        .applymap(color_cell, subset=['全月涨跌幅 (%)', '上半月涨跌幅 (%)', '下半月涨跌幅 (%)'])
        .set_properties(**{'text-align': 'center'})
    )
    
    st.dataframe(styled_df, use_container_width=True, height=400)
    
    # 显示统计指标
    st.subheader("关键统计指标")
    cols = st.columns(4)
    
    with cols[0]:
        avg_change = results_df['全月涨跌幅 (%)'].mean()
        st.metric("平均涨跌幅", f"{avg_change:.2f}%")
    
    with cols[1]:
        positive_months = len(results_df[results_df['全月涨跌幅 (%)'] > 0])
        st.metric("上涨月份", f"{positive_months}次", f"{positive_months/10:.0%}")
    
    with cols[2]:
        max_change = results_df['全月涨跌幅 (%)'].max()
        st.metric("最大涨幅", f"{max_change:.2f}%")
    
    with cols[3]:
        min_change = results_df['全月涨跌幅 (%)'].min()
        st.metric("最大跌幅", f"{min_change:.2f}%")
    
    # 可视化图表
    st.subheader("趋势可视化")
    fig, ax = plt.subplots(figsize=(12, 6))
    sns.lineplot(
        data=results_df,
        x='年份',
        y='全月涨跌幅 (%)',
        marker='o',
        ax=ax
    )
    ax.axhline(0, color='gray', linestyle='--')
    ax.set_title(f"{selected_month}月涨跌幅年度变化")
    ax.set_ylabel("涨跌幅 (%)")
    st.pyplot(fig)

def app():
    df = get_index_data()
    if not df.empty:
        display_analysis(df)
    else:
        st.error("无法加载数据，请检查网络连接")

if __name__ == "__main__":
    app()
